#******************************************************************************
# * Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# *
# * NVIDIA CORPORATION and its licensors retain all intellectual property
# * and proprietary rights in and to this software, related documentation
# * and any modifications thereto.  Any use, reproduction, disclosure or
# * distribution of this software and related documentation without an express
# * license agreement from NVIDIA CORPORATION is strictly prohibited.
# ******************************************************************************

cmake_minimum_required(VERSION 3.18)

project(cuda-matmul-bench LANGUAGES CXX CUDA)

set(CMAKE_CXX_STANDARD           14)
set(CMAKE_CXX_STANDARD_REQUIRED  ON)
set(CMAKE_CXX_EXTENSIONS         OFF)
set(CMAKE_CUDA_STANDARD          14)
set(CMAKE_CUDA_STANDARD_REQUIRED ON)
set(CMAKE_CUDA_EXTENSIONS        OFF)

find_package(OpenMP REQUIRED)
find_package(CUDAToolkit REQUIRED)

if(WIN32)
    set(CUSPARSELT_PATH_DEFAULT "C:/Program Files/NVIDIA cuSPARSELt/v0.7")
else()
    set(CUSPARSELT_PATH_DEFAULT "/usr")
endif()
set(CUSPARSELT_PATH "${CUSPARSELT_PATH_DEFAULT}" CACHE PATH "Path to cuSPARSELt installation")

find_path(CUSPARSELT_INCLUDE_DIR
    cusparseLt.h
    HINTS ${CUSPARSELT_PATH}
    PATH_SUFFIXES include
    DOC "cuSPARSELt include directory"
)
if(NOT CUSPARSELT_INCLUDE_DIR)
    message(FATAL_ERROR "cuSPARSELt headers not found under ${CUSPARSELT_PATH}/include")
endif()

find_library(CUSPARSELT_LIBRARY
    NAMES cusparseLt
    HINTS ${CUSPARSELT_PATH}
    PATH_SUFFIXES lib lib64
    DOC "cuSPARSELt library"
)
if(NOT CUSPARSELT_LIBRARY)
    message(FATAL_ERROR "cuSPARSELt library not found under ${CUSPARSELT_PATH}/lib or lib64")
endif()

add_executable(sparse_fp32 ${PROJECT_SOURCE_DIR}/sparse.cpp)
add_executable(sparse_fp16 ${PROJECT_SOURCE_DIR}/sparse.cpp)
add_executable(sparse_int8 ${PROJECT_SOURCE_DIR}/sparse.cpp)
add_executable(sparse_fp8  ${PROJECT_SOURCE_DIR}/sparse.cpp)
add_executable(sparse_fp4  ${PROJECT_SOURCE_DIR}/sparse.cpp)

target_compile_definitions(sparse_fp32 PRIVATE AB_TYPE=FP32)
target_compile_definitions(sparse_fp16 PRIVATE AB_TYPE=FP16)
target_compile_definitions(sparse_int8 PRIVATE AB_TYPE=INT8)
target_compile_definitions(sparse_fp8  PRIVATE AB_TYPE=FP8 )
target_compile_definitions(sparse_fp4  PRIVATE AB_TYPE=FP4 )

target_include_directories(sparse_fp32 PRIVATE ${CUDAToolkit_INCLUDE_DIRS} ${CUSPARSELT_INCLUDE_DIR})
target_include_directories(sparse_fp16 PRIVATE ${CUDAToolkit_INCLUDE_DIRS} ${CUSPARSELT_INCLUDE_DIR})
target_include_directories(sparse_int8 PRIVATE ${CUDAToolkit_INCLUDE_DIRS} ${CUSPARSELT_INCLUDE_DIR})
target_include_directories(sparse_fp8  PRIVATE ${CUDAToolkit_INCLUDE_DIRS} ${CUSPARSELT_INCLUDE_DIR})
target_include_directories(sparse_fp4  PRIVATE ${CUDAToolkit_INCLUDE_DIRS} ${CUSPARSELT_INCLUDE_DIR})

target_link_libraries(sparse_fp32 PRIVATE OpenMP::OpenMP_CXX CUDA::cudart CUDA::cusparse ${CUSPARSELT_LIBRARY})
target_link_libraries(sparse_fp16 PRIVATE OpenMP::OpenMP_CXX CUDA::cudart CUDA::cusparse ${CUSPARSELT_LIBRARY})
target_link_libraries(sparse_int8 PRIVATE OpenMP::OpenMP_CXX CUDA::cudart CUDA::cusparse ${CUSPARSELT_LIBRARY})
target_link_libraries(sparse_fp8  PRIVATE OpenMP::OpenMP_CXX CUDA::cudart CUDA::cusparse ${CUSPARSELT_LIBRARY})
target_link_libraries(sparse_fp4  PRIVATE OpenMP::OpenMP_CXX CUDA::cudart CUDA::cusparse ${CUSPARSELT_LIBRARY})
